// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <vector>
#include <cstdint>
#include <kllm/proto/interface.struct.pb.h>

namespace kllm {


    struct server_metrics {
        Metrics metrics;

        void init() {
            metrics.set_t_start(ggml_time_us());
        }

        void on_prompt_eval(const server_slot &slot) {
            metrics.set_n_prompt_tokens_processed_total(
                    metrics.n_prompt_tokens_processed_total() + slot.n_prompt_tokens_processed);
            metrics.set_n_prompt_tokens_processed(slot.n_prompt_tokens_processed + metrics.n_prompt_tokens_processed());
            metrics.set_t_prompt_processing(slot.t_prompt_processing + metrics.t_prompt_processing());
            metrics.set_t_prompt_processing_total(slot.t_prompt_processing + metrics.t_prompt_processing_total());
        }

        void on_prediction(const server_slot &slot) {
            metrics.set_n_tokens_predicted_total(slot.n_decoded + metrics.n_tokens_predicted_total());
            metrics.set_n_tokens_predicted(slot.n_decoded + metrics.n_tokens_predicted());
            metrics.set_t_tokens_generation(slot.t_token_generation + metrics.t_tokens_generation());
            metrics.set_t_tokens_generation_total(slot.t_token_generation + metrics.t_tokens_generation_total());
        }

        void on_decoded(const std::vector<server_slot> &slots) {
            metrics.set_n_decode_total(metrics.n_decode_total() + 1);
            int64_t n_busy = 0;
            for (const auto &slot: slots) {
                if (slot.is_processing()) {
                    n_busy++;
                }
            }
            metrics.set_n_busy_slots_total(n_busy);
        }

        void reset_bucket() {
            metrics.set_n_prompt_tokens_processed(0);
            metrics.set_t_prompt_processing(0);
            metrics.set_n_tokens_predicted(0);
            metrics.set_t_tokens_generation(0);
        }
    };


}  // namespace kllm
